// 未完成代码

#include <ros/ros.h>
#include <geometry_msgs/Twist.h>
#include <sensor_msgs/PointCloud2.h>
#include <sensor_msgs/point_cloud2_iterator.h>
#include <tf/transform_broadcaster.h>
#include <tf/transform_listener.h>
#include <geometry_msgs/PoseWithCovarianceStamped.h>
#include <random>

#include <cmath>
#include <mutex>
#include <vector>
#include <algorithm>
#include <numeric>
#include <stdexcept>

class NavSimulate3D
{
public:
    NavSimulate3D(ros::NodeHandle &nh) : nh_(nh), publish_counter_(0),
                                         random_engine_(std::random_device{}())
    {
        // 初始化参数
        nh_.param("initial_x", x_, 0.0);
        nh_.param("initial_y", y_, 0.0);
        nh_.param("initial_z", z_, 2.0);
        nh_.param("initial_yaw", yaw_, 0.0);

        // 物理参数
        nh_.param("gravity", gravity_, -9.81);                           // 重力加速度 (m/s²)
        nh_.param("collision_radius", collision_radius_, 0.2);           // 碰撞检测半径
        nh_.param("bounce_factor", bounce_factor_, 0.3);                 // 碰撞反弹系数
        nh_.param("friction_coeff", friction_coeff_, 0.95);              // 摩擦系数
        nh_.param("pointcloud_noise_stddev", noise_stddev_, 0.005);      // 点云噪声标准差
        nh_.param("max_vertical_velocity", max_vertical_velocity_, 5.0); // 最大下落速度

        vertical_velocity_ = 0.0; // 初始垂直速度

        // 订阅速度指令
        cmd_vel_sub_ = nh_.subscribe("/cmd_vel", 10, &NavSimulate3D::cmdVelCallback, this);

        // 订阅3D点云（使用深拷贝）
        pointcloud_sub_ = nh_.subscribe("/map3d", 10, &NavSimulate3D::pointcloudCallback, this);

        // 发布局部点云
        local_pointcloud_pub_ = nh_.advertise<sensor_msgs::PointCloud2>("/pc_accumulate", 10);

        // 订阅初始位姿
        initialpose_sub_ = nh_.subscribe("/initialpose", 10, &NavSimulate3D::poseResetCallback, this);

        // 定时器设置（20Hz）
        timer_ = nh_.createTimer(ros::Duration(0.05), &NavSimulate3D::timerCallback, this);

        last_time_ = ros::Time::now();

        ROS_INFO("3D Navigation Simulator started with gravity, collision, and pointcloud noise.");
    }

private:
    void cmdVelCallback(const geometry_msgs::TwistConstPtr &msg)
    {
        linear_velocity_ = msg->linear.x;
        angular_velocity_ = msg->angular.z;
    }

    void pointcloudCallback(const sensor_msgs::PointCloud2ConstPtr &msg)
    {
        std::lock_guard<std::mutex> lock(pointcloud_mutex_);

        // 检查是否包含 x/y/z 字段
        bool has_x = false, has_y = false, has_z = false;
        for (const auto &field : msg->fields)
        {
            if (field.name == "x")
                has_x = true;
            else if (field.name == "y")
                has_y = true;
            else if (field.name == "z")
                has_z = true;
        }

        if (!has_x || !has_y || !has_z)
        {
            ROS_ERROR("Received point cloud missing required fields");
            return;
        }

        // 深拷贝点云数据
        global_pointcloud_.reset(new sensor_msgs::PointCloud2(*msg));
        ROS_INFO("Received point cloud with %d points.", msg->width);
    }

    void timerCallback(const ros::TimerEvent &)
    {
        ros::Time current_time = ros::Time::now();
        double dt = (current_time - last_time_).toSec();

        if (dt <= 0)
            return; // 跳过无效时间间隔

        // ====================== 物理模拟 ======================
        // 1. 应用重力
        vertical_velocity_ += gravity_ * dt;

        // 限制最大下落速度（防止穿透）
        vertical_velocity_ = std::max(vertical_velocity_, -max_vertical_velocity_);

        // 2. 更新水平位置
        double dx = linear_velocity_ * cos(yaw_) * dt;
        double dy = linear_velocity_ * sin(yaw_) * dt;
        double prev_x = x_;
        double prev_y = y_;
        double prev_z = z_;

        x_ += dx;
        y_ += dy;
        z_ += vertical_velocity_ * dt;
        yaw_ += angular_velocity_ * dt;

        // 3. 碰撞检测与响应
        bool collision_detected = false;
        if (global_pointcloud_)
        {
            collision_detected = checkCollisions(dt, prev_x, prev_y, prev_z);
        }

        // 4. 地面摩擦
        if (on_ground_ && !collision_detected)
        {
            linear_velocity_ *= friction_coeff_;
            angular_velocity_ *= friction_coeff_;
        }

        // ====================== 发布更新 ======================
        last_time_ = current_time;

        // 发布TF变换（20Hz）
        publishTransform();

        // 发布局部点云（10Hz）
        publish_counter_++;
        if (publish_counter_ % 2 == 0 && global_pointcloud_)
        {
            try
            {
                sensor_msgs::PointCloud2 local_cloud = createLocalPointCloud();
                local_pointcloud_pub_.publish(local_cloud);
            }
            catch (const std::exception &e)
            {
                ROS_ERROR_STREAM("Point cloud processing error: " << e.what());
            }
        }

        // 重置超低高度的机器人
        if (z_ < -100.0)
        {
            resetToOrigin();
        }
    }

    bool checkCollisions(double dt, double prev_x, double prev_y, double prev_z)
    {
        std::lock_guard<std::mutex> lock(pointcloud_mutex_);
        bool collision_detected = false;
        on_ground_ = false;

        // 安全访问点云数据
        sensor_msgs::PointCloud2ConstIterator<float> iter_x(*global_pointcloud_, "x");
        sensor_msgs::PointCloud2ConstIterator<float> iter_y(*global_pointcloud_, "y");
        sensor_msgs::PointCloud2ConstIterator<float> iter_z(*global_pointcloud_, "z");

        // 碰撞检测参数
        const double collision_radius_sq = collision_radius_ * collision_radius_;
        double ground_distance = std::numeric_limits<double>::max();
        double ground_z = z_;

        for (; iter_x != iter_x.end(); ++iter_x, ++iter_y, ++iter_z)
        {
            double dx = *iter_x - x_;
            double dy = *iter_y - y_;
            double dz = *iter_z - z_;
            double dist_sq = dx * dx + dy * dy + dz * dz;

            // 碰撞检测
            if (dist_sq < collision_radius_sq)
            {
                collision_detected = true;

                // 计算碰撞法向量
                double dist = std::sqrt(dist_sq);
                double nx = dx / dist;
                double ny = dy / dist;
                double nz = dz / dist;

                // 计算相对速度在法线方向的分量
                double rel_vel = vertical_velocity_ * nz + linear_velocity_ * (cos(yaw_) * nx + sin(yaw_) * ny);

                // 仅当物体接近点时响应碰撞
                if (rel_vel < 0)
                {
                    // 反弹效果
                    double bounce = -rel_vel * bounce_factor_;

                    // 应用反弹
                    vertical_velocity_ += bounce * nz;
                    linear_velocity_ += bounce * nx;
                    angular_velocity_ += bounce * ny * 0.5; // 旋转效果

                    // 防止穿透
                    double penetration = collision_radius_ - dist;
                    x_ += penetration * nx * 0.5;
                    y_ += penetration * ny * 0.5;
                    z_ += penetration * nz * 0.5;
                }
            }

            // 地面检测（只考虑下方的点）
            if (dz > 0 && dz < 1.0)
            {
                double hor_dist_sq = dx * dx + dy * dy;
                if (hor_dist_sq < 4.0 * collision_radius_sq) // 更大的地面检测范围
                {
                    if (dz < ground_distance)
                    {
                        ground_distance = dz;
                        ground_z = *iter_z;
                    }
                }
            }
        }

        // 地面接触检测
        if (ground_distance < collision_radius_)
        {
            on_ground_ = true;

            // 如果接近地面，应用地面约束
            if (z_ < ground_z + collision_radius_)
            {
                z_ = ground_z + collision_radius_;
                vertical_velocity_ = std::max(0.0, vertical_velocity_); // 防止负速度
            }
        }

        return collision_detected;
    }

    sensor_msgs::PointCloud2 createLocalPointCloud()
    {
        if (!global_pointcloud_ || global_pointcloud_->width == 0)
        {
            ROS_WARN("Global point cloud is empty or not available");
            return sensor_msgs::PointCloud2(); // 返回空点云
        }

        sensor_msgs::PointCloud2 cloud;
        cloud.header.stamp = ros::Time::now();
        cloud.header.frame_id = "body";
        cloud.height = 1;
        cloud.is_dense = true;

        // 正确设置点云字段
        sensor_msgs::PointCloud2Modifier modifier(cloud);
        modifier.setPointCloud2Fields(3,
                                      "x", 1, sensor_msgs::PointField::FLOAT32,
                                      "y", 1, sensor_msgs::PointField::FLOAT32,
                                      "z", 1, sensor_msgs::PointField::FLOAT32);

        // 安全访问点云数据
        sensor_msgs::PointCloud2ConstIterator<float> iter_x(*global_pointcloud_, "x");
        sensor_msgs::PointCloud2ConstIterator<float> iter_y(*global_pointcloud_, "y");
        sensor_msgs::PointCloud2ConstIterator<float> iter_z(*global_pointcloud_, "z");

        double cos_yaw = cos(yaw_);
        double sin_yaw = sin(yaw_);

        // 第一次遍历：统计有效点数量
        size_t valid_points = 0;
        for (; iter_x != iter_x.end(); ++iter_x, ++iter_y, ++iter_z)
        {
            double dx = *iter_x - x_;
            double dy = *iter_y - y_;
            double dz = *iter_z - z_;

            // 转换到局部坐标系
            double x_local = dx * cos_yaw + dy * sin_yaw;
            double y_local = -dx * sin_yaw + dy * cos_yaw;
            double z_local = dz;

            // 裁剪范围
            if (std::abs(x_local) < 5.0 &&
                std::abs(y_local) < 5.0 &&
                std::abs(z_local) < 5.0)
            {
                valid_points++;
            }
        }

        // 重置迭代器
        iter_x = sensor_msgs::PointCloud2ConstIterator<float>(*global_pointcloud_, "x");
        iter_y = sensor_msgs::PointCloud2ConstIterator<float>(*global_pointcloud_, "y");
        iter_z = sensor_msgs::PointCloud2ConstIterator<float>(*global_pointcloud_, "z");

        // 分配精确内存
        modifier.resize(valid_points);

        // 创建写入迭代器
        sensor_msgs::PointCloud2Iterator<float> out_x(cloud, "x");
        sensor_msgs::PointCloud2Iterator<float> out_y(cloud, "y");
        sensor_msgs::PointCloud2Iterator<float> out_z(cloud, "z");

        // 设置随机数生成器（高斯分布）
        std::normal_distribution<float> noise_dist(0.0, noise_stddev_);

        // 第二次遍历：写入有效点并添加噪声
        for (; iter_x != iter_x.end(); ++iter_x, ++iter_y, ++iter_z)
        {
            double dx = *iter_x - x_;
            double dy = *iter_y - y_;
            double dz = *iter_z - z_;

            // 转换到局部坐标系
            double x_local = dx * cos_yaw + dy * sin_yaw;
            double y_local = -dx * sin_yaw + dy * cos_yaw;
            double z_local = dz;

            // 裁剪范围
            if (std::abs(x_local) < 5.0 &&
                std::abs(y_local) < 5.0 &&
                std::abs(z_local) < 5.0)
            {
                // 添加高斯噪声（仅当噪声标准差大于0时）
                if (noise_stddev_ > 0.0)
                {
                    x_local += noise_dist(random_engine_);
                    y_local += noise_dist(random_engine_);
                    z_local += noise_dist(random_engine_);
                }

                *out_x = x_local;
                *out_y = y_local;
                *out_z = z_local;

                // 递增输出迭代器
                ++out_x;
                ++out_y;
                ++out_z;
            }
        }

        cloud.width = valid_points;
        cloud.row_step = cloud.point_step * valid_points;

        return cloud;
    }

    void publishTransform()
    {
        try
        {
            tf::Quaternion q;
            q.setRPY(0, 0, yaw_);

            tf_broadcaster_.sendTransform(tf::StampedTransform(
                tf::Transform(q, tf::Vector3(x_, y_, z_)),
                ros::Time::now(), "map", "body"));
        }
        catch (const std::exception &e)
        {
            ROS_ERROR_STREAM("TF transform error: " << e.what());
        }
    }

    void poseResetCallback(const geometry_msgs::PoseWithCovarianceStampedConstPtr &msg)
    {
        x_ = msg->pose.pose.position.x;
        y_ = msg->pose.pose.position.y;
        z_ = msg->pose.pose.position.z;

        tf::Quaternion q(
            msg->pose.pose.orientation.x,
            msg->pose.pose.orientation.y,
            msg->pose.pose.orientation.z,
            msg->pose.pose.orientation.w);

        double roll, pitch, yaw;
        tf::Matrix3x3(q).getRPY(roll, pitch, yaw);
        yaw_ = yaw;

        // 重置物理状态
        vertical_velocity_ = 0.0;
        linear_velocity_ = 0.0;
        angular_velocity_ = 0.0;

        last_time_ = ros::Time::now();
    }

    void resetToOrigin()
    {
        x_ = 0.0;
        y_ = 0.0;
        z_ = 0.0;
        yaw_ = 0.0;
        vertical_velocity_ = 0.0;
        linear_velocity_ = 0.0;
        angular_velocity_ = 0.0;
        ROS_WARN("Reset to origin due to falling below -100m.");
    }

    ros::NodeHandle nh_;
    ros::Subscriber cmd_vel_sub_;
    ros::Subscriber pointcloud_sub_;
    ros::Publisher local_pointcloud_pub_;
    ros::Subscriber initialpose_sub_;
    ros::Timer timer_;

    // 状态变量
    double x_, y_, z_, yaw_;
    double linear_velocity_, angular_velocity_;
    double vertical_velocity_; // 垂直速度
    ros::Time last_time_;
    sensor_msgs::PointCloud2Ptr global_pointcloud_;
    std::mutex pointcloud_mutex_;
    tf::TransformBroadcaster tf_broadcaster_;
    int publish_counter_;

    // 物理参数
    double gravity_;               // 重力加速度
    double collision_radius_;      // 碰撞检测半径
    double bounce_factor_;         // 碰撞反弹系数
    double friction_coeff_;        // 摩擦系数
    double noise_stddev_;          // 点云噪声标准差
    double max_vertical_velocity_; // 最大垂直下落速度
    bool on_ground_ = false;       // 是否在地面上

    // 随机数生成器
    std::mt19937 random_engine_; // Mersenne Twister 随机引擎
};

int main(int argc, char **argv)
{
    ros::init(argc, argv, "nav_simulate_3d");
    ros::NodeHandle nh;
    NavSimulate3D nav_sim(nh);
    ros::spin();
    return 0;
}